{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results and models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup and load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/arthur/.conda/envs/py38_cu11_2/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "import numpy as np\n",
    "from src.motion_refiner_4D import Motion_refiner, MAX_NUM_OBJS\n",
    "from src.config import *\n",
    "from src.functions import *\n",
    "\n",
    "def evaluate_model(model, x_t, y_t):\n",
    "\n",
    "    print(\"\\nwith next waypoint prediction\")\n",
    "    result_eval = model.evaluate(x_t,y_t)\n",
    "    print(\"MSE: \",result_eval)\n",
    "    print(\"---------------------------------------------------\")\n",
    "    print(\"with autoregressive generation:\")\n",
    "\n",
    "    pred = generate(model ,x_t, traj_n=traj_n).numpy()\n",
    "    result_gen = np.average((y_t - pred[:,1:,:])**2)\n",
    "    print(\"MSE: \",result_gen)\n",
    "    print(\"Trajectory metrics:\")\n",
    "    metrics, metrics_h = compute_metrics(y_t.numpy()[:,:,:],pred[:,1:,:])\n",
    "    return metrics\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### dataset size experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading BERT model... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_layer_norm.weight', 'vocab_transform.weight', 'vocab_transform.bias', 'vocab_projector.bias', 'vocab_layer_norm.bias']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "loading dataset:  latte_100k_lf ...done\n",
      "raw X: (100000, 953) \tY: (100000, 160)\n",
      "filtered X: (96718, 953) \tY: (96718, 160)\n",
      "Train X: (67702, 953) \tY: (67702, 160)\n",
      "Test  X: (19344, 953) \tY: (19344, 160)\n",
      "Val   X: (9672, 953) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=False)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================================================================================\n",
      "Dataset size:  1.0 %\t augmentation:  OFF \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.01, 'augment': 0}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.01-augment:0.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 16s 23ms/step - loss: 9.2614e-04\n",
      "MSE:  0.0009380120318382978\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:26<00:00, 11.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.026081953403375196\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.46488338701590753\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 8.207000090290196\n",
      "mae:\t 0.11063311492115666\n",
      "mse:\t 0.026081953403375113\n",
      "======================================================================================================\n",
      "Dataset size:  1.0 %\t augmentation:  ON  \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.01, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.01-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 16s 22ms/step - loss: 5.0294e-04\n",
      "MSE:  0.00051028624875471\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:18<00:00, 11.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.014203454044915438\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.3573711908340173\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 5.352905033503025\n",
      "mae:\t 0.07590188144170358\n",
      "mse:\t 0.014203454044915383\n",
      "======================================================================================================\n",
      "Dataset size:  10.0 %\t augmentation:  OFF \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.1, 'augment': 0}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.1-augment:0.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 1.3069e-04\n",
      "MSE:  0.00013290416973177344\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:17<00:00, 11.23s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002435481360766444\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.11683638486097525\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.470160787666334\n",
      "mae:\t 0.02347701585593545\n",
      "mse:\t 0.0024354813607664204\n",
      "======================================================================================================\n",
      "Dataset size:  10.0 %\t augmentation:  ON  \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.1, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.1-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 1.3911e-04\n",
      "MSE:  0.00014078985259402543\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:18<00:00, 11.24s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.0024806484219106507\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.1159301913118411\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.508415924637376\n",
      "mae:\t 0.023246562985204867\n",
      "mse:\t 0.002480648421910659\n",
      "======================================================================================================\n",
      "Dataset size:  50.0 %\t augmentation:  OFF \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.5, 'augment': 0}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:0.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 9.0249e-05\n",
      "MSE:  9.204711386701092e-05\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:17<00:00, 11.22s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002241024889381608\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.10978990013544507\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.3282010349645272\n",
      "mae:\t 0.021051997831146235\n",
      "mse:\t 0.002241024889381613\n",
      "======================================================================================================\n",
      "Dataset size:  50.0 %\t augmentation:  ON  \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.5, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 16s 22ms/step - loss: 1.0774e-04\n",
      "MSE:  0.00010925108654191718\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:16<00:00, 11.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.0023629698830414823\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.11344960920497135\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.463039142605034\n",
      "mae:\t 0.022455827967532375\n",
      "mse:\t 0.0023629698830414827\n",
      "======================================================================================================\n",
      "Dataset size:  100.0 %\t augmentation:  OFF \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 0}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:0.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 9.5734e-05\n",
      "MSE:  9.730462625157088e-05\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:16<00:00, 11.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002290856997408395\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.11175631639873772\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.3930190131558344\n",
      "mae:\t 0.022014576386467935\n",
      "mse:\t 0.002290856997408402\n",
      "======================================================================================================\n",
      "Dataset size:  100.0 %\t augmentation:  ON  \n",
      "\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 1.1004e-04\n",
      "MSE:  0.00011141923459945247\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:17<00:00, 11.21s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.00234289053695731\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.11383532855574399\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.469922672372537\n",
      "mae:\t 0.022737672659383878\n",
      "mse:\t 0.0023428905369572905\n"
     ]
    }
   ],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_dataset_size_aug_fixsteps/\"\n",
    "\n",
    "model_names =  [\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.01-augment:0.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.01-augment:1.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.1-augment:0.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.1-augment:1.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:0.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:1.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:0.h5\",\n",
    "                \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\"]\n",
    "models_metrics = {}\n",
    "for model_name in model_names:\n",
    "\n",
    "    model_file = model_path+model_name\n",
    "    param = file_name2dict(model_file,delimiter=\"-\",show=False)\n",
    "    print(\"======================================================================================================\")\n",
    "    print(\"Dataset size: \",str(100.0*float(param[\"sf\"])),\"%\\t\",\"augmentation: \",\"ON \"if param['augment']==1 else \"OFF\", \"\\n\")\n",
    "    model_tag = \"size:\" + str(100.0*float(param[\"sf\"]))+\"_aug:\"+ str(param['augment'])\n",
    "\n",
    "    model = load_model(model_file, delimiter=\"-\")\n",
    "    metrics = evaluate_model(model, x_t, y_t)\n",
    "    models_metrics[model_tag] = metrics\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "========== aug 1 ===========\n",
      "\t1k\t10k\t50k\t100k\t\n",
      "mse\t0.014203454044915383\t0.002480648421910659\t0.0023629698830414827\t0.0023428905369572905\t\n",
      "mae\t0.07590188144170358\t0.023246562985204867\t0.022455827967532375\t0.022737672659383878\t\n",
      "dtw\t5.352905033503025\t2.508415924637376\t2.463039142605034\t2.469922672372537\t\n",
      "dfd\t0.3573711908340173\t0.1159301913118411\t0.11344960920497135\t0.11383532855574399\t\n",
      "========= aug 0 ==========\n",
      "mse\t0.026081953403375113\t0.0024354813607664204\t0.002241024889381613\t0.002290856997408402\t\n",
      "mae\t0.11063311492115666\t0.02347701585593545\t0.021051997831146235\t0.022014576386467935\t\n",
      "dtw\t8.207000090290196\t2.470160787666334\t2.3282010349645272\t2.3930190131558344\t\n",
      "dfd\t0.46488338701590753\t0.11683638486097525\t0.10978990013544507\t0.11175631639873772\t\n"
     ]
    }
   ],
   "source": [
    "print(\"========== aug 1 ===========\\n\\t\",end=\"\")\n",
    "for m in [\"1k\",\"10k\",\"50k\",\"100k\"]:\n",
    "    print(m, end=\"\\t\")\n",
    "print()\n",
    "for k in [\"mse\",\"mae\",\"dtw\",\"dfd\"]:\n",
    "    print(k, end=\"\\t\")\n",
    "    for n,m in models_metrics.items():\n",
    "        if \"aug:1\" in n:\n",
    "            print(m[k], end=\"\\t\")\n",
    "    print()\n",
    "print(\"========= aug 0 ==========\")\n",
    "for k in [\"mse\",\"mae\",\"dtw\",\"dfd\"]:\n",
    "    print(k, end=\"\\t\")\n",
    "    for n,m in models_metrics.items():\n",
    "        if \"aug:0\" in n:\n",
    "            print(m[k], end=\"\\t\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Num of encoders, decoder and model depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================================================================================\n",
      "num Encoders:  1 \tnum Decoders:  3 \tDepth:  256\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 3, 'd_model': 256, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:1-num_layers_dec:3-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 9s 12ms/step - loss: 1.4763e-04\n",
      "MSE:  0.0001494946627644822\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [04:16<00:00,  6.57s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002596976841045097\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.12208694415308749\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.73886212261568\n",
      "mae:\t 0.02640541335892475\n",
      "mse:\t 0.0025969768410450877\n",
      "======================================================================================================\n",
      "num Encoders:  1 \tnum Decoders:  3 \tDepth:  400\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 3, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:1-num_layers_dec:3-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 10s 15ms/step - loss: 1.2021e-04\n",
      "MSE:  0.00012197512842249125\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [04:57<00:00,  7.63s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.0027276886623445857\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.12122696939703666\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.702918660119998\n",
      "mae:\t 0.02529149711046403\n",
      "mse:\t 0.002727688662344588\n",
      "======================================================================================================\n",
      "num Encoders:  1 \tnum Decoders:  5 \tDepth:  256\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 256, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:1-num_layers_dec:5-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 13s 18ms/step - loss: 1.5300e-04\n",
      "MSE:  0.00015453933156095445\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [06:12<00:00,  9.55s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002726531303711237\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.12806367503062693\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.913398924493158\n",
      "mae:\t 0.02838061532357939\n",
      "mse:\t 0.0027265313037112566\n",
      "======================================================================================================\n",
      "num Encoders:  1 \tnum Decoders:  5 \tDepth:  400\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 15s 22ms/step - loss: 9.7113e-05\n",
      "MSE:  9.826025780057535e-05\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:14<00:00, 11.14s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002451821950968988\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.11872049748753903\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.602604161995969\n",
      "mae:\t 0.024537796733149944\n",
      "mse:\t 0.0024518219509689866\n",
      "======================================================================================================\n",
      "num Encoders:  2 \tnum Decoders:  3 \tDepth:  256\n",
      "{'num_layers_enc': 2, 'num_layers_dec': 3, 'd_model': 256, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:2-num_layers_dec:3-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 10s 14ms/step - loss: 1.6059e-04\n",
      "MSE:  0.00016342209710273892\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [05:01<00:00,  7.73s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.0028254162351774624\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.13036128186189053\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.981241223477885\n",
      "mae:\t 0.030203137270130145\n",
      "mse:\t 0.0028254162351774585\n",
      "======================================================================================================\n",
      "num Encoders:  2 \tnum Decoders:  3 \tDepth:  400\n",
      "{'num_layers_enc': 2, 'num_layers_dec': 3, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:2-num_layers_dec:3-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 12s 17ms/step - loss: 1.2579e-04\n",
      "MSE:  0.00012719212099909782\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [05:50<00:00,  8.99s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.002568349239366221\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.12146609359756542\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.7679229428692076\n",
      "mae:\t 0.027026672595851592\n",
      "mse:\t 0.0025683492393662063\n",
      "======================================================================================================\n",
      "num Encoders:  2 \tnum Decoders:  5 \tDepth:  256\n",
      "{'num_layers_enc': 2, 'num_layers_dec': 5, 'd_model': 256, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:2-num_layers_dec:5-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 14s 20ms/step - loss: 1.6047e-04\n",
      "MSE:  0.00016268411127384752\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [06:49<00:00, 10.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.00282736868302177\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.13193735104718837\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 3.0200743898124935\n",
      "mae:\t 0.03050950657869114\n",
      "mse:\t 0.002827368683021769\n",
      "======================================================================================================\n",
      "num Encoders:  2 \tnum Decoders:  5 \tDepth:  400\n",
      "{'num_layers_enc': 2, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 1.0, 'augment': 1}\n",
      "loading weights:  /home/azureuser/data/models/ICRA_TF4D_enc_dec_depth/TF-num_layers_enc:2-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "605/605 [==============================] - 17s 24ms/step - loss: 9.9642e-05\n",
      "MSE:  0.00010097025369759649\n",
      "---------------------------------------------------\n",
      "with autoregressive generation:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [08:02<00:00, 12.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MSE:  0.0026205959150691696\n",
      "Trajectory metrics:\n",
      "pcm:\t 0.0\n",
      "dfd:\t 0.1242314966561472\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.7840958375592493\n",
      "mae:\t 0.02682884721046181\n",
      "mse:\t 0.0026205959150691523\n"
     ]
    }
   ],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_enc_dec_depth/\"\n",
    "\n",
    "model_names =  [\n",
    "\"TF-num_layers_enc:1-num_layers_dec:3-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:1-num_layers_dec:3-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:1-num_layers_dec:5-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:2-num_layers_dec:3-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:2-num_layers_dec:3-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:2-num_layers_dec:5-d_model:256-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\",\n",
    "\"TF-num_layers_enc:2-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\"\n",
    "]\n",
    "models_metrics_enc_dec = {}\n",
    "for model_name in model_names:\n",
    "\n",
    "    model_file = model_path+model_name\n",
    "    param = file_name2dict(model_file,delimiter=\"-\",show=False)\n",
    "    print(\"======================================================================================================\")\n",
    "    print(\"num Encoders: \",param[\"num_layers_enc\"],\"\\tnum Decoders: \",param[\"num_layers_dec\"],\"\\tDepth: \",param['d_model'])\n",
    "    model_tag = \"enc:\"+str(param[\"num_layers_enc\"])+\"-dec:\"+str(param[\"num_layers_dec\"])+\"-d\"+str(param['d_model'])\n",
    "\n",
    "    model = load_model(model_file, delimiter=\"-\")\n",
    "    metrics = evaluate_model(model, x_t, y_t)\n",
    "    models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse\tmae\tdtw\tdfd\t\n",
      "enc:1-dec:3-d256\t0.0025969768410450877\t0.02640541335892475\t2.73886212261568\t0.12208694415308749\t\n",
      "enc:1-dec:3-d400\t0.002727688662344588\t0.02529149711046403\t2.702918660119998\t0.12122696939703666\t\n",
      "enc:1-dec:5-d256\t0.0027265313037112566\t0.02838061532357939\t2.913398924493158\t0.12806367503062693\t\n",
      "enc:1-dec:5-d400\t0.0024518219509689866\t0.024537796733149944\t2.602604161995969\t0.11872049748753903\t\n",
      "enc:2-dec:3-d256\t0.0028254162351774585\t0.030203137270130145\t2.981241223477885\t0.13036128186189053\t\n",
      "enc:2-dec:3-d400\t0.0025683492393662063\t0.027026672595851592\t2.7679229428692076\t0.12146609359756542\t\n",
      "enc:2-dec:5-d256\t0.002827368683021769\t0.03050950657869114\t3.0200743898124935\t0.13193735104718837\t\n",
      "enc:2-dec:5-d400\t0.0026205959150691523\t0.02682884721046181\t2.7840958375592493\t0.1242314966561472\t\n"
     ]
    }
   ],
   "source": [
    "for k in [\"mse\",\"mae\",\"dtw\",\"dfd\"]:\n",
    "    print(k, end=\"\\t\")\n",
    "print()\n",
    "for n,m in models_metrics.items():\n",
    "    if \"dec\" in n:\n",
    "        print(n, end=\"\\t\")\n",
    "        for k in [\"mse\",\"mae\",\"dtw\",\"dfd\"]:\n",
    "            print(m[k], end=\"\\t\")\n",
    "        print()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Naive predictor\n",
    "Coping initial trajectory (no modifications)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading BERT model... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_transform.bias', 'vocab_layer_norm.bias', 'vocab_transform.weight', 'vocab_layer_norm.weight', 'vocab_projector.bias']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "loading dataset:  latte_100k_lf ...done\n",
      "raw X: (100000, 953) \tY: (100000, 160)\n",
      "filtered X: (96718, 953) \tY: (96718, 160)\n",
      "Train X: (67702, 953) \tY: (67702, 160)\n",
      "Test  X: (19344, 953) \tY: (19344, 160)\n",
      "Val   X: (9672, 953) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=False)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "19344it [08:34, 37.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pcm:\t 0.0\n",
      "dfd:\t 0.13866536138405858\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 3.5680424145512766\n",
      "mae:\t 0.027094356069811483\n",
      "mse:\t 0.004370182933815549\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model_tag=\"naive_predictor\"\n",
    "traj_in = x_t[0][:,MAX_NUM_OBJS+1:,:].numpy()\n",
    "traj_out = y_t.numpy()\n",
    "metrics, metrics_h = compute_metrics(traj_out,traj_in)\n",
    "models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### No language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "clip only\n",
      "loading dataset:  latte_100k_lf_cliponly ...done\n",
      "raw X: (100000, 697) \tY: (100000, 160)\n",
      "filtered X: (96718, 697) \tY: (96718, 160)\n",
      "Train X: (67702, 697) \tY: (67702, 160)\n",
      "Test  X: (19344, 697) \tY: (19344, 160)\n",
      "Val   X: (9672, 697) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=True)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf_cliponly\", filter_data = True, base_path=data_folder)\n",
    "X[:,feature_indices] = 0\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_no_language/\"\n",
    "\n",
    "model_name=  \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:537-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\"\n",
    "model_tag=\"no_language\"\n",
    "model_file = model_path+model_name\n",
    "model = load_model(model_file, delimiter=\"-\")\n",
    "metrics = evaluate_model(model, x_t, y_t)\n",
    "models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 39/39 [07:18<00:00, 11.25s/it]\n"
     ]
    }
   ],
   "source": [
    "pred = generate(model ,x_t, traj_n=traj_n).numpy()\n",
    "data_pred = np.array(data)[indices_test].copy()\n",
    "for i,d in enumerate(data_pred):\n",
    "    d[\"output_traj\"] = pred[i]\n",
    "\n",
    "mr.save_data(data_pred,data_name=\"test_no_language_100k_latte_f\", base_path=data_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder Only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading BERT model... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_projector.bias', 'vocab_transform.bias', 'vocab_transform.weight', 'vocab_layer_norm.weight', 'vocab_layer_norm.bias']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "loading dataset:  latte_100k_lf ...done\n",
      "raw X: (100000, 953) \tY: (100000, 160)\n",
      "filtered X: (96718, 953) \tY: (96718, 160)\n",
      "Train X: (67702, 953) \tY: (67702, 160)\n",
      "Test  X: (19344, 953) \tY: (19344, 160)\n",
      "Val   X: (9672, 953) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=False)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 16, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': True, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse'}\n",
      "loading weights:  /home/azureuser/data/models/FINAL_decoder_only/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:16-bs:16-dense_n:512-num_dense:3-concat_emb:True-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse.h5\n",
      "\n",
      "with next waypoint prediction\n",
      "---------------------------------------------------\n",
      "Trajectory metrics:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "19344it [09:06, 35.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pcm:\t 0.0\n",
      "dfd:\t 0.12620982754684149\n",
      "area:\t 0.0\n",
      "cl:\t 0.0\n",
      "dtw:\t 2.290391868346655\n",
      "mae:\t 0.02421539465031163\n",
      "mse:\t 0.0016075996776673247\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'models_metrics' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_1903644/3192676914.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_file\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdelimiter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"-\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0mmetrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate_model_dec_only\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx_t_deconly\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_t\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m \u001b[0mmodels_metrics\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmodel_tag\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmetrics\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'models_metrics' is not defined"
     ]
    }
   ],
   "source": [
    "from src.TF4D_decoder_only import *\n",
    "model_path = models_folder+\"FINAL_decoder_only/\"\n",
    "\n",
    "model_name=  \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:16-bs:16-dense_n:512-num_dense:3-concat_emb:True-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse.h5\"\n",
    "\n",
    "def evaluate_model_dec_only(model, x_t, y_t):\n",
    "\n",
    "    print(\"\\nwith next waypoint prediction\")\n",
    "    # result_eval = model.evaluate(x_t,y_t)\n",
    "    # print(\"MSE: \",result_eval)\n",
    "    print(\"---------------------------------------------------\")\n",
    "    pred = model.predict(x_t)\n",
    "    print(\"Trajectory metrics:\")\n",
    "    metrics, metrics_h = compute_metrics(y_t.numpy()[:,:,:],pred[:,:,:])\n",
    "    return metrics\n",
    "\n",
    "x_t_deconly = (x_t[0][:,MAX_NUM_OBJS:-1,:],x_t[1],x_t[2])\n",
    "model_tag=\"decoder_only\"\n",
    "model_file = model_path+model_name\n",
    "model = load_model(model_file, delimiter=\"-\")\n",
    "metrics = evaluate_model_dec_only(model,x_t_deconly, y_t)\n",
    "models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model.predict(x_t_deconly)\n",
    "data_pred = np.array(data)[indices_test].copy()\n",
    "for i,d in enumerate(data_pred):\n",
    "    d[\"output_traj\"] = pred[i]\n",
    "\n",
    "mr.save_data(data_pred,data_name=\"test_decoder_only_100k_latte_f\", base_path=data_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### clip only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "clip only\n",
      "loading dataset:  latte_100k_lf_cliponly ...done\n",
      "raw X: (100000, 697) \tY: (100000, 160)\n",
      "filtered X: (96718, 697) \tY: (96718, 160)\n",
      "Train X: (67702, 697) \tY: (67702, 160)\n",
      "Test  X: (19344, 697) \tY: (19344, 160)\n",
      "Val   X: (9672, 697) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "import numpy as np\n",
    "from src.motion_refiner_4D import Motion_refiner, MAX_NUM_OBJS\n",
    "from src.config import *\n",
    "from src.functions import *\n",
    "\n",
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=True)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf_cliponly\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_clip_only/\"\n",
    "\n",
    "model_name=  \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:537-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\"\n",
    "model_tag=\"decoder_only\"\n",
    "model_file = model_path+model_name\n",
    "model = load_model(model_file, delimiter=\"-\")\n",
    "metrics = evaluate_model(model, x_t, y_t)\n",
    "models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text Only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading BERT model... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_layer_norm.bias', 'vocab_transform.weight', 'vocab_layer_norm.weight', 'vocab_projector.bias', 'vocab_projector.weight', 'vocab_transform.bias']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "loading dataset:  latte_100k_lf_textonly ...done\n",
      "raw X: (100000, 953) \tY: (100000, 160)\n",
      "filtered X: (96718, 953) \tY: (96718, 160)\n",
      "Train X: (67702, 953) \tY: (67702, 160)\n",
      "Test  X: (19344, 953) \tY: (19344, 160)\n",
      "Val   X: (9672, 953) \tY: (9672, 160)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "import numpy as np\n",
    "from src.motion_refiner_4D import Motion_refiner, MAX_NUM_OBJS\n",
    "from src.config import *\n",
    "from src.functions import *\n",
    "\n",
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=False)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"latte_100k_lf_textonly\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((mr.prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_text_only/\"\n",
    "\n",
    "model_name=  \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:537-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:1.0-augment:1.h5\"\n",
    "model_tag=\"decoder_only\"\n",
    "model_file = model_path+model_name\n",
    "model = load_model(model_file, delimiter=\"-\")\n",
    "metrics = evaluate_model(model, x_t, y_t)\n",
    "models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forces interactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading BERT model... "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_transform.bias', 'vocab_transform.weight', 'vocab_projector.weight', 'vocab_layer_norm.bias', 'vocab_projector.bias', 'vocab_layer_norm.weight']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n",
      "loading CLIP model... done\n",
      "DEVICE:  cuda\n",
      "loading dataset:  forces_only_f ...done\n",
      "raw X: (10000, 953) \tY: (10000, 160)\n",
      "filtered X: (10000, 953) \tY: (10000, 160)\n",
      "Train X: (7000, 953) \tY: (7000, 160)\n",
      "Test  X: (2000, 953) \tY: (2000, 160)\n",
      "Val   X: (1000, 953) \tY: (1000, 160)\n",
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 16, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': True, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse'}\n",
      "loading weights:  /home/azureuser/data/models/forces_onl/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:16-bs:16-dense_n:512-num_dense:3-concat_emb:True-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse.h5\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "import numpy as np\n",
    "from src.motion_refiner_4D import Motion_refiner, MAX_NUM_OBJS\n",
    "from src.config import *\n",
    "from src.functions import *\n",
    "\n",
    "mr = Motion_refiner(load_models=True ,traj_n = traj_n, locality_factor=True, clip_only=False)\n",
    "feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()\n",
    "embedding_indices = mr.embedding_indices\n",
    "\n",
    "#============================== load dataset ==========================================\n",
    "X,Y, data = mr.load_dataset(\"forces_only_f\", filter_data = True, base_path=data_folder)\n",
    "X_train, X_test, X_valid, y_train, y_test, y_valid, indices_train, indices_test, indices_val = mr.split_dataset(X, Y, test_size=0.2, val_size=0.1)\n",
    "\n",
    "\n",
    "def prepare_x(x):\n",
    "    objs = pad_array(list_to_wp_seq(x[:,obj_poses_indices],d=3),4,axis=-1) # no speed\n",
    "    trajs = list_to_wp_seq(x[:,traj_indices],d=4)\n",
    "    #   return np.concatenate([objs,trajs],axis = 1)\n",
    "    return trajs[:,:-1,:]\n",
    "\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((prepare_x(X_test),\n",
    "                                                    list_to_wp_seq(y_test,d=4),\n",
    "                                                    X_test[:,embedding_indices])).batch(X_test.shape[0])\n",
    "\n",
    "g = generator(test_dataset,stop=True,augment=False)\n",
    "x_t, y_t = next(g)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 16, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': True, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse'}\n",
      "loading weights:  /home/azureuser/data/models/forces_onl/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:16-bs:16-dense_n:512-num_dense:3-concat_emb:True-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse.h5\n"
     ]
    }
   ],
   "source": [
    "from src.TF4D_decoder_only import *\n",
    "\n",
    "model_path = models_folder+\"forces_onl/\"\n",
    "model_name = \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:16-bs:16-dense_n:512-num_dense:3-concat_emb:True-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse.h5\"\n",
    "\n",
    "model_file = model_path+model_name\n",
    "model = load_model(model_file, delimiter=\"-\")\n",
    "\n",
    "# metrics = evaluate_model(model, x_t, y_t)\n",
    "# models_metrics[model_tag] = metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "63/63 [==============================] - 1s 19ms/step\n"
     ]
    }
   ],
   "source": [
    "pred = model.predict(x_t)\n",
    "pred_d = np.array(data)[indices_test]\n",
    "for i,d in enumerate(pred_d):\n",
    "    d[\"forces\"] = pred[i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "indices = np.random.choice(range(len(indices_test)), 3)\n",
    "\n",
    "plt.close('all')\n",
    "# pred_t = np.transpose(pred[:,:,:2],[0,2,1])\n",
    "data_array = np.array(data)[indices_test[indices]]\n",
    "show_data4D(data_array,plot_output=False)\n",
    "# data_array = pred_d[indices]\n",
    "# show_data4D(data_array,plot_output=False)\n",
    "\n",
    "# show_data4D(data_array,plot_output=False,image_loader=mr.image_loader, color_traj=False, change_img_base=[\"/home/mirmi/Arthur/dataset/\",\"/home/tum/data/image_dataset/\"])\n",
    "# show_data4D(data_sample,plot_output=False)#,image_loader=mr.image_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'num_layers_enc': 1, 'num_layers_dec': 5, 'd_model': 400, 'dff': 512, 'num_heads': 8, 'dropout_rate': 0.1, 'wp_d': 4, 'num_emb_vec': 4, 'bs': 16, 'dense_n': 512, 'num_dense': 3, 'concat_emb': False, 'features_n': 793, 'optimizer': 'adam', 'norm_layer': True, 'activation': 'tanh', 'loss': 'mse', 'sf': 0.5, 'augment': 1}\n",
      "loading weights:  /home/arthur/local_data/models/FINAL_dataset_size_aug_fixsteps/TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:1.h5\n"
     ]
    }
   ],
   "source": [
    "from src.TF4D_mult_features import *\n",
    "model_path = models_folder+\"FINAL_dataset_size_aug_fixsteps/\"\n",
    "\n",
    "model_name = \"TF-num_layers_enc:1-num_layers_dec:5-d_model:400-dff:512-num_heads:8-dropout_rate:0.1-wp_d:4-num_emb_vec:4-bs:16-dense_n:512-num_dense:3-concat_emb:False-features_n:793-optimizer:adam-norm_layer:True-activation:tanh-loss:mse-sf:0.5-augment:1.h5\"\n",
    "model_file = model_path+model_name\n",
    "\n",
    "model = load_model(model_file, delimiter=\"-\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = generate(model ,x_t, traj_n=traj_n).numpy()\n",
    "data_pred = np.array(data)[indices_test].copy()\n",
    "for i,d in enumerate(data_pred):\n",
    "    d[\"output_traj\"] = pred[i]\n",
    "\n",
    "mr.save_data(data_pred,data_name=\"testpred_100k_latte_f\", base_path=data_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# indices = np.random.choice(range(len(indices_test)), 3)\n",
    "indices = np.arange(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_pred = mr.load_data(\"testpred_100k_latte_f\", base_path=data_folder)\n",
    "pred = np.array([d[\"output_traj\"] for d in data_pred])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from src.functions import *\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "# plot best results\n",
    "th = 0.0005 # maximum mse\n",
    "good_indices = {\"dist\":[], \"cartesian\":[], \"speed\":[]}\n",
    "for i,d in enumerate(np.array(data)[indices_test]):\n",
    "    mse = np.mean((pred[i]-np.array(d[\"output_traj\"]))**2)\n",
    "    if mse < th and len(d[\"obj_names\"])<4:\n",
    "        good_indices[d[\"change_type\"]] = good_indices[d[\"change_type\"]] + [i]\n",
    "indices = [np.random.choice(good_indices[i]) for i in good_indices.keys()]\n",
    "indices[0] =6276#11699#,11664\n",
    "# indices[1] =3190\n",
    "# indices[2] = 14168#11053, 16036, 14168\n",
    "indices = indices\n",
    "print(indices)\n",
    "\n",
    "data_array = np.array(data)[indices_test[indices]]\n",
    "\n",
    "show_data4D(data_array, pred=pred[indices,:,:],color_traj=False,labels=[\"Ground Truth\",\"Predicted\"],image_loader=mr.image_loader, change_img_base=[\"/mnt/tumdata/image_dataset/\", image_dataset_folder])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('py38_cu11_2')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5de91200c9f9e1f8a0c28ceba668014be0fd55838e84400e0a7ad1d269192773"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
